module ub_atomic_rdgen (
    input logic clk,
    input logic rst_n,
    aru_ub_wr_req_if.in u_aru_wr_req_in_if,
    gdma_ub_wr_req_if.out u_aru_wr_req_to_fifo_if,
    gdma_ub_wr_req_if.out u_aru_wr_req_to_wr_arbiter_if,
    ub_rd_req_if.out u_ub_atomic_rd_req_if
);

    // Pipeline register signals
    logic                                     reg_vld;
    logic                                     reg_atomic_mode;
    logic [                             31:0] reg_ub_addr;
    logic [$bits(u_aru_wr_req_in_if.pld)-1:0] reg_pld;

    // Pipeline control
    logic                                     pipe_enable;
    logic                                     output_rdy;

    // Determine if pipeline can accept new data
    assign pipe_enable            = (!reg_vld) || output_rdy;
    assign u_aru_wr_req_in_if.rdy = pipe_enable;

    // Pipeline register update
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            reg_vld         <= 1'b0;
            reg_atomic_mode <= 1'b0;
            reg_ub_addr     <= '0;
            reg_pld         <= '0;
        end else begin
            if (pipe_enable) begin
                reg_vld <= u_aru_wr_req_in_if.vld;
                if (u_aru_wr_req_in_if.vld) begin
                    reg_atomic_mode <= (u_aru_wr_req_in_if.pld.atomic_mode != 2'b00);
                    reg_ub_addr     <= u_aru_wr_req_in_if.pld.ub_addr;
                    reg_pld         <= u_aru_wr_req_in_if.pld;
                end
            end
        end
    end

    // Helper function: expand ARU 64B write request to GDMA 128B format
    // Lower 64B: from ARU request, Upper 64B: mask set to 0
    function automatic gdma_ub_wr_req_s expand_aru_to_gdma(aru_ub_wr_req_s aru_req);
        gdma_ub_wr_req_s gdma_req;

        // Address remains the same
        gdma_req.ub_addr     = aru_req.ub_addr;

        // Copy atomic_mode
        gdma_req.atomic_mode = aru_req.atomic_mode;

        // Lower 64B (first P_ARU*N0 elements): copy from ARU request
        for (int i = 0; i < `P_ARU * `N0; i++) begin
            gdma_req.dat[i] = aru_req.dat[i];
            gdma_req.msk[i] = aru_req.msk[i];
        end

        // Upper 64B (next P_ARU*N0 elements): mask = 0 (invalid)
        for (int i = `P_ARU * `N0; i < `N0 * `K0; i++) begin
            gdma_req.dat[i] = '0;
            gdma_req.msk[i] = 1'b0;  // Mask = 0 means don't write
        end

        return gdma_req;
    endfunction

    // For atomic mode: generate read request and send write request to FIFO
    // For normal mode: directly forward write request to WrArbiter
    always_comb begin
        if (reg_atomic_mode) begin
            // Atomic mode: generate read request
            u_ub_atomic_rd_req_if.vld         = reg_vld;
            u_ub_atomic_rd_req_if.ub_addr     = reg_ub_addr;
            u_ub_atomic_rd_req_if.src         = 2'b01;  // atomic read

            // Send write request to FIFO (expand 64B -> 128B)
            u_aru_wr_req_to_fifo_if.vld       = reg_vld;
            u_aru_wr_req_to_fifo_if.pld       = expand_aru_to_gdma(reg_pld);

            // No direct write to WrArbiter in atomic mode
            u_aru_wr_req_to_wr_arbiter_if.vld = 1'b0;
            u_aru_wr_req_to_wr_arbiter_if.pld = '0;

            // Backpressure: ready when both atomic read and FIFO can accept
            output_rdy                        = u_ub_atomic_rd_req_if.rdy && u_aru_wr_req_to_fifo_if.rdy;

        end else begin
            // Normal write mode: forward directly to WrArbiter (expand 64B -> 128B)
            u_aru_wr_req_to_wr_arbiter_if.vld = reg_vld;
            u_aru_wr_req_to_wr_arbiter_if.pld = expand_aru_to_gdma(reg_pld);

            // No atomic read request
            u_ub_atomic_rd_req_if.vld         = 1'b0;
            u_ub_atomic_rd_req_if.ub_addr     = '0;
            u_ub_atomic_rd_req_if.src         = 2'b00;

            // No write to FIFO in normal mode
            u_aru_wr_req_to_fifo_if.vld       = 1'b0;
            u_aru_wr_req_to_fifo_if.pld       = '0;

            // Backpressure: ready when WrArbiter can accept
            output_rdy                        = u_aru_wr_req_to_wr_arbiter_if.rdy;
        end
    end

endmodule
